# -*- coding: UTF-8 -*- 
import sys
sys.path.append('/home/xd233/.local/lib/python3.6/site-packages')
import cv2
import time
import socket

import matplotlib.pyplot as plt
import nanocamera as nano
from utils import *
import os
import numpy as np
import glob
import pickle
from map2array_lib import map2array, array2senddata, generate_start_end,verify_map,modify_map
from neural_net_classify.conv_net import ConvNet

import serial
from uart_example import get_uart_data
import RPi.GPIO as GPIO
led_pin = 19
button_pin=31
high_pin=33
beep_pin=29
GPIO.setmode(GPIO.BOARD)
GPIO.setup(led_pin, GPIO.OUT)
GPIO.setup(beep_pin, GPIO.OUT,initial=GPIO.LOW)
def beep_t(t):
    GPIO.output(beep_pin, GPIO.HIGH)
    s=time.perf_counter()
    while(time.perf_counter()-s)<t:
        pass
    GPIO.output(beep_pin, GPIO.LOW)
GPIO.setup(high_pin, GPIO.OUT)
GPIO.output(high_pin, GPIO.HIGH)
GPIO.setup(button_pin, GPIO.IN, pull_up_down=GPIO.PUD_UP)
ttt=my_timer()
time.sleep(0.01)
GPIO.output(led_pin, GPIO.HIGH)
time.sleep(0.3)
GPIO.output(led_pin, GPIO.LOW)
time.sleep(0.3)
GPIO.output(led_pin, GPIO.HIGH)
time.sleep(0.3)
GPIO.output(led_pin, GPIO.LOW)
time.sleep(0.3)
GPIO.output(led_pin, GPIO.HIGH)
time.sleep(0.3)
GPIO.output(led_pin, GPIO.LOW)
time.sleep(0.3)
GPIO.output(led_pin, GPIO.HIGH)
time.sleep(0.3)
if GPIO.input(button_pin):
    GPIO.output(led_pin, GPIO.HIGH)
    group_flag=1
    print("we are red")
else:
    GPIO.output(led_pin, GPIO.LOW)
    group_flag=0
    print("we are blue")

width = 1280
height = 720
camera_1 = nano.Camera(device_id=1, flip=2, width=width, height=height, fps=60)#,enforce_fps=True)  #

camera_2 = nano.Camera(device_id=0, flip=0, width=width, height=height, fps=10)#,enforce_fps=True)  #
with open('mtx.pickle', "rb") as f1:
    mtx = pickle.load(f1)
with open('dist.pickle', "rb") as f1:
    dist = pickle.load(f1)

h, w = height, width
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))

# undistort
mapx, mapy = cv2.initUndistortRectifyMap(mtx, dist, None, newcameramtx, (w, h), 5)

serial_port = serial.Serial(
    port="/dev/ttyTHS1",
    baudrate=115200,
    bytesize=serial.EIGHTBITS,
    parity=serial.PARITY_NONE,
    stopbits=serial.STOPBITS_ONE,
)

def ten2byteofhex(num):
    hex_str = str(hex(num))[2:]
    if len(hex_str) == 1:
        hex_str = '0' + hex_str
    return (bytes.fromhex(hex_str))
def refresh(data_array,end_p,r_or_f,obj_num):
    data_array[0,end_p[0],end_p[1]]=0
   # if r_or_f ==1:
   #     data_array[0,9-end_p[0],9-end_p[1]]=0
   # elif r_or_f ==0:
   #     if obj_num==1 or obj_num==2:#为伪宝藏
   #         data_array[0,9-end_p[0],9-end_p[1]]=0
    return data_array


def judge_and_send(obj_num):  # {0:"blue_yellow",1:"blue_green",2:"red_yellow",3:"red_green"}
    print(obj_num)            #0,3为对应真宝藏，1，2为伪宝藏
    global camera_modle
    global group_flag
    global end_p
    global data_array
    global pass_treasure_num
    global pass_real_treasure_num
    global real_treasure
    pass_treasure_num = pass_treasure_num + 1
    if group_flag==0:
        real_treasure=0
    else:
        real_treasure=3
    if obj_num == real_treasure:  # 为真宝藏
        serial_port.write(bytes.fromhex("AA0301FF"))
        camera_modle=2.2
        print("real\n")
        pass_real_treasure_num=pass_real_treasure_num+1
        data_array=refresh(data_array,end_p,1,obj_num)
    else:  # 为假宝藏
        serial_port.write(bytes.fromhex("AA0302FF"))
        print("fake\n")
        camera_modle=1
        data_array=refresh(data_array,end_p,0,obj_num)
        read_pipline(camera_1,0.06)

    with open('./pause_cache/data_array.pickle', "wb") as f2:
        pickle.dump(data_array, f2)
    with open('./pause_cache/num.pickle', "wb") as f3:
        pickle.dump([pass_treasure_num,pass_real_treasure_num], f3)

mask_track = np.zeros([int(300 * 0.3), int(1280 * 0.3)]).astype("uint8")
(h, w) = mask_track.shape
for i in range(h):
    for j in range(w):
        if (250 * 0.3) ** 2 < (i - (h - 1)) ** 2 + ((j - (w // 2)) / 1.6) ** 2 < (260 * 0.3) ** 2:
            mask_track[i, j] = 1

testnet = ConvNet()
testnet.load_params("./neural_net_classify/out1.pkl")

camera_modle = 0  # 0什么都不做   1发送巡线信息   2发送宝藏位置
per_pc = 1
cn = 0
T = 1 / 30
x=0

def modify_map_read_gpio(data_array):
    data_array[0,9,0]=0
    data_array[0,0,9]=0
    GPIO.setup(36, GPIO.OUT)#36,38为高电平引脚
    GPIO.output(36, GPIO.HIGH)
    GPIO.setup(38, GPIO.OUT)
    GPIO.output(38, GPIO.HIGH)

    GPIO.setup(35, GPIO.OUT)#35，37为低电平引脚
    GPIO.output(35, GPIO.LOW)
    GPIO.setup(37, GPIO.OUT)
    GPIO.output(37, GPIO.LOW)

    GPIO.setup(24, GPIO.IN)
    GPIO.setup(23, GPIO.IN)
    if GPIO.input(23):#23号引脚对应左下角，若为高(接36或38)，则上方有墙,若为低(接35，37)，则右侧有墙
        data_array[1, 9, 0] = 1
        data_array[3, 8, 0] = 1

        data_array[4, 9, 0] = 0
        data_array[2, 9, 1] = 0
        print("左下上方有墙")
    else:
        data_array[1, 9, 0] = 0
        data_array[3, 8, 0] = 0

        data_array[4, 9, 0] = 1
        data_array[2, 9, 1] = 1
        print("左下右方有墙")
    if GPIO.input(24):#24号引脚对应右上角，若为高(接36或38)，则下方有墙,若为低(接35，37)，则左侧有墙
        data_array[3, 0, 9] = 1
        data_array[1, 1, 9] = 1

        data_array[2, 0, 9] = 0
        data_array[4, 0, 8] = 0
        print("右上下方有墙")
    else:
        data_array[3, 0, 9] = 0
        data_array[1, 1, 9] = 0

        data_array[2, 0, 9] = 1
        data_array[4, 0, 8] = 1
        print("右上左方有墙")
    return data_array

x=192
while True:
   # try:
    cn = cn + 1
    ts = time.perf_counter()
    uart_data = get_uart_data()  # jetson nano 接收到的数据，根据\n来分段

    if camera_modle == 1:
        cv_image = camera_1.read()
        last_x=x
        x, y = track_line(cv_image, send_data[index_senddata][0], mask_track)
        if x==None:
            x=last_x
        serial_port.write(bytes.fromhex("AA01") + str(x).encode() + bytes.fromhex("FF"))
      #  print("direction:"+str(send_data[index_senddata][0])+",num:"+str(x))
    elif camera_modle == 2.1:
        serial_port.write(bytes.fromhex("AA01") + str(192).encode() + bytes.fromhex("FF"))
        cv_image = camera_1.read()
        dst = cv2.remap(cv_image, mapx, mapy, cv2.INTER_LINEAR)
        # crop the image
        x, y, w, h = roi
        cv_image = dst[y:y + h, x:x + w]
        crop_treasure, [y_min, y_max, x_min, x_max] = crop_treasure_picture(cv_image)
        print(y_min)
        if crop_treasure is not None:
            if crop_treasure[0][0][0].dtype == 'uint8':
                crop_treasure = crop_treasure / 255
            input_img = cv2.resize(crop_treasure, (30, 60))
            input_img = np.expand_dims(input_img, axis=0)
            input_img = input_img.transpose(0, 3, 1, 2)
            out = testnet.predict(input_img, train_flg=False)
            obj_num = np.argmax(out)
            if last_obj_num == obj_num:
                detect_num = detect_num + 1
                if detect_num == 2:
                    judge_and_send(obj_num)
                    uart_data = ""
                    detect_num = 0
            else:
                detect_num = 0
            last_obj_num = obj_num
    elif camera_modle == 2.2:
        cv_image = camera_2.read()
        dst = cv2.remap(cv_image, mapx, mapy, cv2.INTER_LINEAR)
        # crop the image
        x, y, w, h = roi
        cv_image = dst[y:y + h, x:x + w]
        crop_treasure, [y_min, y_max, x_min, x_max] = crop_treasure_picture(cv_image)
        if crop_treasure is not None:
           # serial_port.write(bytes.fromhex("AA01") + str(int((x_min+x_max)/2/1025*384)).encode() + bytes.fromhex("FF"))
            print(int((x_min+x_max)/2/1025*384),y_min)
            if y_min>266:
                serial_port.write(bytes.fromhex("AA0303FF"))
                print("掉头")
                camera_modle = 999
                time.sleep(3)
                serial_port.write(bytes.fromhex("AA0304FF"))
    elif camera_modle==0:
        pass

    if len(uart_data) > 1:
        print(uart_data)
        if uart_data == "capture\r\n":
            treasure_num = 8
            detect_num = 0
            last_obj_num = -1
            index_senddata = 1
            end_p = [9, 0]
            uart_data = ""
            camera_modle=0
            if len(os.listdir("./pause_cache"))>0:
                with open('./pause_cache/data_array.pickle', "rb") as f2:
                    data_array = pickle.load(f2)
                with open('./pause_cache/num.pickle', "rb") as f3:
                    num_list = pickle.load(f3)
                #data_array=modify_map_read_gpio(data_array)
                pass_treasure_num = num_list[0]
                pass_real_treasure_num=num_list[1]
                camera_modle = 1  # 进入巡线模式
                start_p = end_p
                end_p = generate_start_end(start_p, data_array)
                send_data = array2senddata(data_array, start_p, end_p)  # 生成要发送给stm32的数据
                data_num = len(send_data)
                serial_port.write(bytes.fromhex("AA02") + ten2byteofhex(data_num) + ten2byteofhex(0))
                for i in range(data_num):
                    serial_port.write(ten2byteofhex(send_data[i][0]))
                    serial_port.write(ten2byteofhex(send_data[i][1]))
                serial_port.write(bytes.fromhex("FF"))
                print("path:\n")
                print(send_data)
                print("\n")
            else:
                read_pipline(camera_2,0.15)
                cv_image = camera_2.read()
                dst = cv2.remap(cv_image, mapx, mapy, cv2.INTER_LINEAR)
                # crop the image
                x, y, w, h = roi
                cv_image = dst[y:y + h, x:x + w]
                pass_treasure_num = 0
                pass_real_treasure_num=0
                data_array, crop_map = map2array(cv_image)  # 迷宫对应的数组数据
                if data_array is not None:
                    pass
                    #data_array=modify_map_read_gpio(data_array)
                if verify_map(data_array,treasure_num,0)==0:
                    data_array = None
                    print("capture map is not right")
                if data_array is not None:
                    #data_array=modify_map(data_array)
                    with open('./pause_cache/data_array.pickle', "wb") as f2:
                        pickle.dump(data_array, f2)
                    with open('./pause_cache/num.pickle', "wb") as f3:
                        pickle.dump([pass_treasure_num,pass_real_treasure_num], f3)

                    camera_modle = 1  # 进入巡线模式
                    print("map:\n")
                    print(data_array[0, :, :])
                    print("down:")
                    print(data_array[3, :, :])
                    print("left:")
                    print(data_array[2, :, :])
                    start_p = end_p
                    end_p = generate_start_end(start_p, data_array)
                    print("start and end:\n")
                    print(start_p, end_p)

                    send_data = array2senddata(data_array, start_p, end_p)  # 生成要发送给stm32的数据
                    data_num = len(send_data)
                    serial_port.write(bytes.fromhex("AA02") + ten2byteofhex(data_num) + ten2byteofhex(0))
                    for i in range(data_num):
                        serial_port.write(ten2byteofhex(send_data[i][0]))
                        serial_port.write(ten2byteofhex(send_data[i][1]))
                    serial_port.write(bytes.fromhex("FF"))
                    print("path:\n")
                    print(send_data)
                    print("\n")
                else:
                    print("capture is error")
        elif uart_data == "detect\r\n":
            ttt.tic()
            read_pipline(camera_2,0.06)
            ttt.toc()
            camera_modle = 2.1
            if end_p==[0,9]:
                camera_modle=999
        elif uart_data == "next\r\n":
            read_pipline(camera_1,0.06)
            uart_data = ""
            index_senddata = 1  # 还原路书索引为1
            camera_modle = 1  # 进入巡线模式
            start_p = end_p
            if pass_treasure_num == treasure_num or pass_real_treasure_num==3:
                end_p = [0, 9]
            else:
                end_p = generate_start_end(start_p, data_array)
            print(start_p, end_p)
            send_data = array2senddata(data_array, start_p, end_p)  # 生成要发送给stm32的数据
            send_data[0][1] = 0
            data_num = len(send_data)
            serial_port.write(bytes.fromhex("AA02") + ten2byteofhex(data_num) + ten2byteofhex(0))
            for i in range(data_num):
                serial_port.write(ten2byteofhex(send_data[i][0]))
                serial_port.write(ten2byteofhex(send_data[i][1]))
            serial_port.write(bytes.fromhex("FF"))
            print(send_data)
        elif uart_data == "index_add\r\n":
            print("get_add and now is "+ str(index_senddata+1))
            uart_data = ""
            if index_senddata == len(send_data) - 1:
                pass
            else:
                index_senddata = index_senddata + 1
  #  except KeyboardInterrupt:
  #      print("key board")
  #      break
  #  except:
  #      print("something is error,try again")

    te = time.perf_counter()
    T = T * 0.9 + 0.1 * (te - ts)
    if cn == per_pc:
        cn = 0
        #print("x:" + str(x) + ",f:" + str(int(1 / T)))
